package persistence

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"sort"
	"strings"
)

// Migration 表示一个数据库迁移
type Migration struct {
	Version int
	Up      string
	Down    string
}

// Migrator 数据库迁移工具
type Migrator struct {
	db         *sql.DB
	migrations []Migration
}

// NewMigrator 创建迁移工具
func NewMigrator(db *sql.DB, migrationsPath string) (*Migrator, error) {
	m := &Migrator{
		db: db,
	}

	// 创建迁移表
	if err := m.createMigrationsTable(); err != nil {
		return nil, err
	}

	// 加载迁移文件
	if err := m.loadMigrations(migrationsPath); err != nil {
		return nil, err
	}

	return m, nil
}

// createMigrationsTable 创建迁移表
func (m *Migrator) createMigrationsTable() error {
	query := `
		CREATE TABLE IF NOT EXISTS migrations (
			version INT PRIMARY KEY,
			applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
		)
	`
	_, err := m.db.Exec(query)
	return err
}

// loadMigrations 加载迁移文件
func (m *Migrator) loadMigrations(migrationsPath string) error {
	files, err := ioutil.ReadDir(migrationsPath)
	if err != nil {
		return err
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}

		if !strings.HasSuffix(file.Name(), ".up.sql") {
			continue
		}

		version := parseVersion(file.Name())
		if version == 0 {
			continue
		}

		upPath := filepath.Join(migrationsPath, file.Name())
		downPath := filepath.Join(migrationsPath, strings.Replace(file.Name(), ".up.sql", ".down.sql", 1))

		upSQL, err := ioutil.ReadFile(upPath)
		if err != nil {
			return err
		}

		downSQL, err := ioutil.ReadFile(downPath)
		if err != nil {
			return err
		}

		m.migrations = append(m.migrations, Migration{
			Version: version,
			Up:      string(upSQL),
			Down:    string(downSQL),
		})
	}

	sort.Slice(m.migrations, func(i, j int) bool {
		return m.migrations[i].Version < m.migrations[j].Version
	})

	return nil
}

// parseVersion 解析版本号
func parseVersion(filename string) int {
	var version int
	fmt.Sscanf(filename, "%d_", &version)
	return version
}

// Migrate 执行迁移
func (m *Migrator) Migrate() error {
	// 获取已应用的迁移
	applied, err := m.getAppliedMigrations()
	if err != nil {
		return err
	}

	// 执行未应用的迁移
	for _, migration := range m.migrations {
		if applied[migration.Version] {
			continue
		}

		if err := m.applyMigration(migration); err != nil {
			return err
		}
	}

	return nil
}

// getAppliedMigrations 获取已应用的迁移
func (m *Migrator) getAppliedMigrations() (map[int]bool, error) {
	query := `SELECT version FROM migrations`
	rows, err := m.db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	applied := make(map[int]bool)
	for rows.Next() {
		var version int
		if err := rows.Scan(&version); err != nil {
			return nil, err
		}
		applied[version] = true
	}

	return applied, nil
}

// applyMigration 应用迁移
func (m *Migrator) applyMigration(migration Migration) error {
	tx, err := m.db.Begin()
	if err != nil {
		return err
	}

	if _, err := tx.Exec(migration.Up); err != nil {
		tx.Rollback()
		return err
	}

	if _, err := tx.Exec("INSERT INTO migrations (version) VALUES (?)", migration.Version); err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}

// Rollback 回滚迁移
func (m *Migrator) Rollback() error {
	// 获取已应用的迁移
	applied, err := m.getAppliedMigrations()
	if err != nil {
		return err
	}

	// 按版本号倒序执行回滚
	for i := len(m.migrations) - 1; i >= 0; i-- {
		migration := m.migrations[i]
		if !applied[migration.Version] {
			continue
		}

		if err := m.rollbackMigration(migration); err != nil {
			return err
		}
	}

	return nil
}

// rollbackMigration 回滚迁移
func (m *Migrator) rollbackMigration(migration Migration) error {
	tx, err := m.db.Begin()
	if err != nil {
		return err
	}

	if _, err := tx.Exec(migration.Down); err != nil {
		tx.Rollback()
		return err
	}

	if _, err := tx.Exec("DELETE FROM migrations WHERE version = ?", migration.Version); err != nil {
		tx.Rollback()
		return err
	}

	return tx.Commit()
}
